#pragma once
#include <iostream>
#include <cassert>

namespace ns_AVL
{
    template <class T>
    struct TreeNode
    {
    public:
        T _key;
        int _bf;            //平衡因子
        TreeNode<T>* _left;
        TreeNode<T>* _right;
        TreeNode<T>* _parent;

        TreeNode(const T& key = T()) : _key(key), _bf(0), _left(nullptr), _right(nullptr), _parent(nullptr) 
        {}
    };

    template <class T>
    class AVL
    {
    public:
        typedef TreeNode<T> Node;

        AVL() : _root(nullptr)
        {}

        ~AVL(){}

        bool insert(const T& key)
        {
            if(_root == nullptr)
            {
                _root = new Node(key);
                return true;
            }
            Node* cur = _root;
            Node* parent = nullptr;
            while(cur)
            {
                parent = cur;
                if(cur->_key > key)
                    cur = cur->_left;
                else if(cur->_key < key)
                    cur = cur->_right;
                else
                    return false;   //节点已经存在了
            }
            cur = new Node(key);
            if(parent->_key > key)
                parent->_left = cur;
            else
                parent->_right = cur;
            cur->_parent = parent;

            //平衡调整
            while(parent)
            {
                if(parent->_left == cur) //调节平衡因子
                    --parent->_bf;
                else
                    ++parent->_bf;

                if(parent->_bf == 0)
                {
                    break;
                }
                else if(parent->_bf == -1 || parent->_bf == 1)
                {
                    cur = parent;
                    parent = parent->_parent;
                }
                else if(cur->_bf == -1 && parent->_bf == -2) //右单旋
                {
                    RotateR(parent);
                    break;
                }   
                else if(cur->_bf == 1 && parent->_bf == -2)  
                {
                    RotateLR(parent);
                    break;
                } 
                else if(cur->_bf == 1 && parent->_bf == 2)  //左单旋
                {
                    RotateL(parent);
                    break;
                }
                else if(cur->_bf == -1 && parent->_bf == 2)
                {
                    RotateRL(parent);
                    break;
                }
                else
                {
                    assert(false);
                }
            }
            return true;
        }

        void Inorder()
        {
            _Inorder(_root);
            std::cout << std::endl;
        }

        bool Balance()
        {
            return _Balance(_root);

        }
    private:
        bool _Balance(Node* root)
        {
            if(root == nullptr)
                return true;
            int leftHeight = Height(root->_left);
            int rightHeight = Height(root->_right);
            if((rightHeight - leftHeight) != root->_bf)
            {
                std::cout << "应该是：" << rightHeight - leftHeight << std::endl;
                std::cout << "实际是：" << root->_bf << std::endl;
                return false;
            }

            return abs(leftHeight-rightHeight) < 2 && _Balance(root->_left) && _Balance(root->_right);
        }

        int Height(Node* root)
        {
            if(nullptr == root)
                return 0;
            int leftHeight = Height(root->_left);
            int rightHeight = Height(root->_right);

            return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
        }

        void _Inorder(Node*& root)
        {
            if(nullptr == root)
                return;
            _Inorder(root->_left);
            std::cout << root->_key << " ";
            _Inorder(root->_right);
        }

        void RotateR(Node*& root)
        {
            Node* Left = root->_left;
            Node* LeftR = Left->_right;

            root->_left = LeftR;
            if(LeftR)
                LeftR->_parent = root;

            Left->_right = root;
            if(root->_parent == nullptr)
            {
                _root = Left;
                _root->_parent = nullptr;
            }
            else
            {
                if(root->_parent->_left == root)
                    root->_parent->_left = Left;
                else
                    root->_parent->_right = Left;
                Left->_parent = root->_parent;
            }
            root->_parent = Left;
            Left->_bf = root->_bf = 0;
        }

        void RotateL(Node*& root)
        {
            Node* Right = root->_right;
            Node* RightL = Right->_left;

            root->_right = RightL;
            if(RightL)
                RightL->_parent = root;
            
            Right->_left = root;
            if(root->_parent == nullptr)
            {
                _root = Right;
                _root->_parent = nullptr;
            }
            else
            {
                if(root->_parent->_left == root)
                    root->_parent->_left = Right;
                else
                    root->_parent->_right = Right;
                Right->_parent = root->_parent;
            }
            root->_parent = Right;
            Right->_bf = root->_bf = 0;
        }

        void RotateLR(Node*& root)
        {
            Node* Left = root->_left;
            Node* LeftR = Left->_right;
            int bf = LeftR->_bf;

            RotateL(Left);
            RotateR(root);
            if(bf == 0)
            {
                root->_bf = 0;
                Left->_bf = 0;
                LeftR->_bf = 0;
            }
            else if(bf == 1)
            {
                root->_bf = 0;
                LeftR->_bf = 0;
                Left->_bf = -1;
            }
            else if(bf == -1)
            {
                root->_bf = 1;
                LeftR->_bf = 0;
                Left->_bf = 0;
            }
            else
            {
                assert(false);
            }
        }

        void RotateRL(Node*& root)
        {
            Node* Right = root->_right;
            Node* RightL = Right->_left;
            int bf = RightL->_bf;

            RotateR(Right);
            RotateL(root);

            if(bf == 0)
            {
                root->_bf = 0;
                Right->_bf = 0;
                RightL->_bf = 0;
            }
            else if(bf == 1)
            {
                root->_bf = -1;
                Right->_bf = 0;
                RightL->_bf = 0;
            }
            else if(bf == -1)
            {
                root->_bf = 0;
                Right->_bf = 1;
                RightL->_bf = 0;
            }
            else
            {
                assert(false);
            }
        }
    private:
        Node* _root;
    };
}